import torch
import torchvision
import torchvision.transforms as transforms
from torch.utils.data import DataLoader
import torch.nn as nn

# 模型参数44426
# 在笔记本上训练minist大概3分钟
# 准确率98.93%

# 超参数设置
input_size = 784     # 输入层维度（28x28像素展平为784）
hidden_size = 512    # 第一个隐藏层神经元数量
num_classes = 10     # 输出类别数（0-9数字）
num_epochs = 5       # 训练轮次
batch_size = 100     # 每批数据量
learning_rate = 0.001  # 学习率

# MNIST数据集加载
def load_data():
    """
    加载MNIST数据集并进行预处理
    返回:
        train_loader: 训练集数据加载器
        test_loader: 测试集数据加载器
    """
    # 训练集配置
    train_dataset = torchvision.datasets.MNIST(
        root='./data',
        train=True,
        transform=transforms.Compose([
            transforms.ToTensor(),  # 将PIL图像转换为Tensor[0,1]
            transforms.Normalize((0.1307,), (0.3081,))  # 标准化（均值, 标准差）
        ]),
        download=True  # 自动下载数据集
    )

    # 测试集配置（使用相同的归一化参数）
    test_dataset = torchvision.datasets.MNIST(
        root='./data',
        train=False,
        transform=transforms.Compose([
            transforms.ToTensor(),
            transforms.Normalize((0.1307,), (0.3081,))
        ])
    )

    # 创建数据加载器
    return (
        DataLoader(train_dataset, batch_size=batch_size, shuffle=True),  # 训练集打乱顺序
        DataLoader(test_dataset, batch_size=batch_size, shuffle=False)    # 测试集保持顺序
    )

# 定义LeNet-5模型
# 详细讲解课参考： https://blog.csdn.net/fly_wt/article/details/95599187
# LeNet-5共有7层，不包含输入，每层都包含可训练参数
class LeNet5(nn.Module):
    def __init__(self):
        super(LeNet5, self).__init__()
        # 卷积层：包含两个卷积层。输入图像大小为 28×28×1，
        #   第一层卷积层使用 5×5 的卷积核，输出 6 个特征图；输出特征图大小为 24×24×6。
        self.conv1 = nn.Conv2d(1, 6, kernel_size=5)
        #   池化层输出12x12x6
        self.pool1 = nn.MaxPool2d(kernel_size=2)
        #   第二层卷积层使用 5×5 的卷积核，输入通道为 6，输出通道为 16。输出特征图大小为 8×8×16。
        self.conv2 = nn.Conv2d(6, 16, kernel_size=5)
        # 池化层：两个池化层分别接在两个卷积层之后，采用 2×2 的池化窗口，进行下采样操作。
        # 输出特征图大小为 4×4×16。
        self.pool2 = nn.MaxPool2d(kernel_size=2)
        # 全连接层：包含三个全连接层。
        #   第一层全连接层将 16×4×4 的特征图展平，并输出为 120 个节点；
        #   第二层全连接层将 120 个节点映射为 84 个节点；
        #   第三层全连接层将 84 个节点映射为 10 个节点，对应 10 个类别。
        #   最后一个全连接层根据具体任务的类别数确定神经元数量，
        #   在 MNIST 数字识别任务中为 10 个，对应 0-9 这 10 个数字类别。
        self.fc1 = nn.Linear(16 * 4 * 4, 120)
        self.fc2 = nn.Linear(120, 84)
        self.fc3 = nn.Linear(84, num_classes)

    def forward(self, x):
        x = self.pool1(torch.relu(self.conv1(x)))
        x = self.pool2(torch.relu(self.conv2(x)))
        x = x.view(-1, 16 * 4 * 4)
        x = torch.relu(self.fc1(x))
        x = torch.relu(self.fc2(x))
        x = self.fc3(x)
        return x

# 训练模型函数
def train_model(model, criterion, optimizer, train_loader):
    model.train()
    for epoch in range(num_epochs):
        running_loss = 0.0
        for i, (images, labels) in enumerate(train_loader):
            optimizer.zero_grad()
            outputs = model(images)
            loss = criterion(outputs, labels)
            loss.backward()
            optimizer.step()
            running_loss += loss.item()

        print(f'Epoch {epoch + 1}/{num_epochs}, Loss: {running_loss / len(train_loader)}')

# 评估模型函数
def evaluate_model(model, test_loader):
    model.eval()
    correct = 0
    total = 0
    with torch.no_grad():
        for images, labels in test_loader:
            outputs = model(images)
            _, predicted = torch.max(outputs.data, 1)
            total += labels.size(0)
            correct += (predicted == labels).sum().item()

    print(f'Accuracy on test set: {100 * correct / total}%')

# 计算模型参数数量的函数
def count_parameters(model):
    return sum(p.numel() for p in model.parameters() if p.requires_grad)

# 主程序入口
if __name__ == '__main__':
    # 数据加载 → 模型初始化 → 训练 → 评估 完整流程
    # 1、数据加载
    train_loader, test_loader = load_data()
    # 2、模型实例化，使用LeNet-5
    model = LeNet5()
    # 打印模型参数数量
    total_params = count_parameters(model)
    print(f'Total number of trainable parameters: {total_params}')
    # 3、损失函数
    criterion = nn.CrossEntropyLoss()  # 交叉熵损失，相当于Softmax+Log+NllLoss
    # 4、优化器
    optimizer = torch.optim.Adam(model.parameters(), lr=learning_rate)  # 使用Adam优化器
    # 5、训练与评估
    train_model(model, criterion, optimizer, train_loader)
    # 6、测试模型
    evaluate_model(model, test_loader)
